{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 诗歌生成"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 数据处理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/remote-home/jjgong01/a3/envs/tf2.0/lib/python3.6/site-packages/h5py/__init__.py:36: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.\n",
      "  from ._conv import register_converters as _register_converters\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import collections\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers\n",
    "from tensorflow.keras import layers, optimizers, datasets\n",
    "\n",
    "start_token = 'bos'\n",
    "end_token = 'eos'\n",
    "\n",
    "def process_dataset(fileName):\n",
    "    examples = []\n",
    "    with open(fileName, 'r') as fd:\n",
    "        for line in fd:\n",
    "            outs = line.strip().split(':')\n",
    "            content = ''.join(outs[1:])\n",
    "            ins = [start_token] + list(content) + [end_token] \n",
    "            if len(ins) > 200:\n",
    "                continue\n",
    "            examples.append(ins)\n",
    "            \n",
    "    counter = collections.Counter()\n",
    "    for e in examples:\n",
    "        for w in e:\n",
    "            counter[w]+=1\n",
    "    \n",
    "    sorted_counter = sorted(counter.items(), key=lambda x: -x[1])  # 排序\n",
    "    words, _ = zip(*sorted_counter)\n",
    "    words = ('PAD', 'UNK') + words[:len(words)]\n",
    "    word2id = dict(zip(words, range(len(words))))\n",
    "    id2word = {word2id[k]:k for k in word2id}\n",
    "    \n",
    "    indexed_examples = [[word2id[w] for w in poem]\n",
    "                        for poem in examples]\n",
    "    seqlen = [len(e) for e in indexed_examples]\n",
    "    \n",
    "    instances = list(zip(indexed_examples, seqlen))\n",
    "    \n",
    "    return instances, word2id, id2word\n",
    "\n",
    "def poem_dataset():\n",
    "    instances, word2id, id2word = process_dataset('../poems.txt')\n",
    "    ds = tf.data.Dataset.from_generator(lambda: [ins for ins in instances], \n",
    "                                            (tf.int64, tf.int64), \n",
    "                                            (tf.TensorShape([None]),tf.TensorShape([])))\n",
    "    ds = ds.shuffle(buffer_size=10240)\n",
    "    ds = ds.padded_batch(100, padded_shapes=(tf.TensorShape([None]),tf.TensorShape([])))\n",
    "    ds = ds.map(lambda x, seqlen: (x[:, :-1], x[:, 1:], seqlen-1))\n",
    "    return ds, word2id, id2word"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 模型代码， 完成建模代码"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class myRNNModel(keras.Model):\n",
    "    def __init__(self, w2id):\n",
    "        super(myRNNModel, self).__init__()\n",
    "        self.v_sz = len(w2id)\n",
    "        self.embed_layer = tf.keras.layers.Embedding(self.v_sz, 64, \n",
    "                                                    batch_input_shape=[None, None])\n",
    "        \n",
    "        self.rnncell = tf.keras.layers.SimpleRNNCell(128)\n",
    "        self.rnn_layer = tf.keras.layers.RNN(self.rnncell, return_sequences=True)\n",
    "        self.dense = tf.keras.layers.Dense(self.v_sz)\n",
    "        \n",
    "    @tf.function\n",
    "    def call(self, inp_ids):\n",
    "        '''\n",
    "        此处完成建模过程，可以参考Learn2Carry\n",
    "        '''\n",
    "        return logits\n",
    "    \n",
    "    @tf.function\n",
    "    def get_next_token(self, x, state):\n",
    "        '''\n",
    "        shape(x) = [b_sz,] \n",
    "        '''\n",
    "    \n",
    "        inp_emb = self.embed_layer(x) #shape(b_sz, emb_sz)\n",
    "        h, state = self.rnncell.call(inp_emb, state) # shape(b_sz, h_sz)\n",
    "        logits = self.dense(h) # shape(b_sz, v_sz)\n",
    "        out = tf.argmax(logits, axis=-1)\n",
    "        return out, state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 一个计算sequence loss的辅助函数，只需了解用途。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mkMask(input_tensor, maxLen):\n",
    "    shape_of_input = tf.shape(input_tensor)\n",
    "    shape_of_output = tf.concat(axis=0, values=[shape_of_input, [maxLen]])\n",
    "\n",
    "    oneDtensor = tf.reshape(input_tensor, shape=(-1,))\n",
    "    flat_mask = tf.sequence_mask(oneDtensor, maxlen=maxLen)\n",
    "    return tf.reshape(flat_mask, shape_of_output)\n",
    "\n",
    "\n",
    "def reduce_avg(reduce_target, lengths, dim):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        reduce_target : shape(d_0, d_1,..,d_dim, .., d_k)\n",
    "        lengths : shape(d0, .., d_(dim-1))\n",
    "        dim : which dimension to average, should be a python number\n",
    "    \"\"\"\n",
    "    shape_of_lengths = lengths.get_shape()\n",
    "    shape_of_target = reduce_target.get_shape()\n",
    "    if len(shape_of_lengths) != dim:\n",
    "        raise ValueError(('Second input tensor should be rank %d, ' +\n",
    "                         'while it got rank %d') % (dim, len(shape_of_lengths)))\n",
    "    if len(shape_of_target) < dim+1 :\n",
    "        raise ValueError(('First input tensor should be at least rank %d, ' +\n",
    "                         'while it got rank %d') % (dim+1, len(shape_of_target)))\n",
    "\n",
    "    rank_diff = len(shape_of_target) - len(shape_of_lengths) - 1\n",
    "    mxlen = tf.shape(reduce_target)[dim]\n",
    "    mask = mkMask(lengths, mxlen)\n",
    "    if rank_diff!=0:\n",
    "        len_shape = tf.concat(axis=0, values=[tf.shape(lengths), [1]*rank_diff])\n",
    "        mask_shape = tf.concat(axis=0, values=[tf.shape(mask), [1]*rank_diff])\n",
    "    else:\n",
    "        len_shape = tf.shape(lengths)\n",
    "        mask_shape = tf.shape(mask)\n",
    "    lengths_reshape = tf.reshape(lengths, shape=len_shape)\n",
    "    mask = tf.reshape(mask, shape=mask_shape)\n",
    "\n",
    "    mask_target = reduce_target * tf.cast(mask, dtype=reduce_target.dtype)\n",
    "\n",
    "    red_sum = tf.reduce_sum(mask_target, axis=[dim], keepdims=False)\n",
    "    red_avg = red_sum / (tf.cast(lengths_reshape, dtype=tf.float32) + 1e-30)\n",
    "    return red_avg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 定义loss函数，定义训练函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def compute_loss(logits, labels, seqlen):\n",
    "    losses = tf.nn.sparse_softmax_cross_entropy_with_logits(\n",
    "            logits=logits, labels=labels)\n",
    "    losses = reduce_avg(losses, seqlen, dim=1)\n",
    "    return tf.reduce_mean(losses)\n",
    "\n",
    "@tf.function\n",
    "def train_one_step(model, optimizer, x, y, seqlen):\n",
    "    '''\n",
    "    完成一步优化过程，可以参考之前做过的模型\n",
    "    '''\n",
    "    return loss\n",
    "\n",
    "def train(epoch, model, optimizer, ds):\n",
    "    loss = 0.0\n",
    "    accuracy = 0.0\n",
    "    for step, (x, y, seqlen) in enumerate(ds):\n",
    "        loss = train_one_step(model, optimizer, x, y, seqlen)\n",
    "\n",
    "        if step % 500 == 0:\n",
    "            print('epoch', epoch, ': loss', loss.numpy())\n",
    "\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 训练优化过程"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: Logging before flag parsing goes to stderr.\n",
      "W0921 02:36:07.816946 140256959960832 deprecation.py:323] From /remote-home/jjgong01/a3/envs/tf2.0/lib/python3.6/site-packages/tensorflow/python/data/ops/dataset_ops.py:410: py_func (from tensorflow.python.ops.script_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "tf.py_func is deprecated in TF V2. Instead, there are two\n",
      "    options available in V2.\n",
      "    - tf.py_function takes a python function which manipulates tf eager\n",
      "    tensors instead of numpy arrays. It's easy to convert a tf eager tensor to\n",
      "    an ndarray (just call tensor.numpy()) but having access to eager tensors\n",
      "    means `tf.py_function`s can use accelerators such as GPUs as well as\n",
      "    being differentiable using a gradient tape.\n",
      "    - tf.numpy_function maintains the semantics of the deprecated tf.py_func\n",
      "    (it is not differentiable, and manipulates numpy arrays). It drops the\n",
      "    stateful argument making all functions stateful.\n",
      "    \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 0 : loss 8.8213\n",
      "epoch 1 : loss 6.508053\n",
      "epoch 2 : loss 5.8712797\n",
      "epoch 3 : loss 5.679646\n",
      "epoch 4 : loss 5.510748\n",
      "epoch 5 : loss 5.367529\n",
      "epoch 6 : loss 5.2667193\n",
      "epoch 7 : loss 5.172158\n",
      "epoch 8 : loss 5.096925\n",
      "epoch 9 : loss 5.0376143\n"
     ]
    }
   ],
   "source": [
    "optimizer = optimizers.Adam(0.0005)\n",
    "train_ds, word2id, id2word = poem_dataset()\n",
    "model = myRNNModel(word2id)\n",
    "\n",
    "for epoch in range(10):\n",
    "    loss = train(epoch, model, optimizer, train_ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 生成过程"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "一年一里，花落花声。eosPAD山风雨，风吹落月。eosPAD山风雨，风吹落月。eosPAD山风雨，风吹落月。eosPAD山风雨，风\n"
     ]
    }
   ],
   "source": [
    "def gen_sentence():\n",
    "    state = [tf.random.normal(shape=(1, 128), stddev=0.5), tf.random.normal(shape=(1, 128), stddev=0.5)]\n",
    "    cur_token = tf.constant([word2id['bos']], dtype=tf.int32)\n",
    "    collect = []\n",
    "    for _ in range(50):\n",
    "        cur_token, state = model.get_next_token(cur_token, state)\n",
    "        collect.append(cur_token.numpy()[0])\n",
    "    return [id2word[t] for t in collect]\n",
    "print(''.join(gen_sentence()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
